# _*_ coding: utf-8 _*_
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, randint, shuffle, uniform
from time import time, gmtime, strftime, localtime
import argparse

from VCMTSP import Instance


np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)



m = 12  # number of salesman, indexes ranging from 0, ..., m-1
case = 'eil51'

instance = Instance(file=f'./TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
Distance = instance.distance_matrix  # distance matrix
Cost = instance.cost  # traveling cost of a single path
A = instance.A  # accessibility matrix
D = instance.depot  # depot
Cities = deepcopy(instance.cities[1:]).tolist()  # city set, exclude depot 1


def make_dir(path):
    path = path.strip()
    path = path.rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        return True
    else:
        print(path + ' already exists.')
        return False


def find_nearest(city, candidates):
    tar, min_dis = -1, 1e99
    for c in candidates:
        if min_dis > Distance[city][c]:
            min_dis = Distance[city][c]
            tar = c
    assert tar != -1
    return tar


def salesmen_of(city):
    return np.where(A[city, :] == 1)[0].tolist()


def cities_of(salesman):
    return np.where(A[:, salesman] == 1)[0].tolist()


class Solution:
    def __init__(self, init='RND'):
        self.cities = None  # city chromosome
        self.salesmen = None  # salesman chromosome
        self.length = None  # total length
        self.lengths = None  # length of each route
        # initialize by random
        if init == 'EMP':
            # empty route
            self.cities = []
            self.salesmen = []
            return
        elif init == 'RND':
            self.cities = deepcopy(Cities)
            shuffle(self.cities)
            self.salesmen = []
            for city in self.cities:  # random appoint salesmen
                self.salesmen.append(sample(salesmen_of(city), 1)[0])
        elif init == 'GRD':
            self.cities = deepcopy(Cities)
            shuffle(self.cities)
            self.salesmen = []
            for city in self.cities:
                self.salesmen.append(sample(salesmen_of(city), 1)[0])
            cities = [[] for _ in range(m)]  # record random assigned cities of each salesman
            for i, city in enumerate(self.cities):
                cities[self.salesmen[i]].append(city)
            pre = [D for _ in range(m)]
            for i, salesman in enumerate(self.salesmen):
                self.cities[i] = find_nearest(city=pre[salesman], candidates=cities[salesman])
                pre[salesman] = self.cities[i]
                cities[salesman].remove(self.cities[i])
        else:
            print(f'Error. init={init}')
            quit(-1)
        self.validate()
        self.calculate_length()

    def calculate_length(self):
        self.lengths = [0 for _ in range(m)]
        pre = [D for _ in range(m)]
        cur = [D for _ in range(m)]
        for i, city in enumerate(self.cities):
            _m = self.salesmen[i]
            cur[_m] = city
            self.lengths[_m] += Distance[pre[_m]][cur[_m]]
            pre[_m] = city
        for _m in range(m):
            self.lengths[_m] += Distance[pre[_m]][D]
        self.length = sum(self.lengths)

    def salesman_correction(self):
        for x in range(0, len(self.cities)):
            if A[self.cities[x]][self.salesmen[x]] == 0:  # current salesman invalid
                salesman = sample(salesmen_of(self.cities[x]), 1)[0]  # random sample a salesman
                self.salesmen[x] = salesman  # random assign another feasible salesman
        return self

    def validate(self):
        diff1 = set(Cities) - set(self.cities)
        diff2 = set(self.cities) - set(Cities)
        assert len(diff1) == 0 & len(diff2) == 0
        for i, city in enumerate(self.cities):
            if A[city][self.salesmen[i]] != 1:
                print(f'Error: city {city} assigned to {self.salesmen[i]}. ')

    def to_routes(self, single=False):
        routes = [[D] for _ in range(m)]
        for i, city in enumerate(self.cities):
            routes[self.salesmen[i]].append(self.cities[i])
        for _ in range(m):
            routes[_].append(D)
        if single:
            res = []
            for route in routes:
                res.extend(route)
            return res
        return routes


def partially_mapped_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = list(range(0, len(P1.cities)))
    a, b = sorted(sample(indexes, 2))
    for i in range(a, b + 1):
        O1.cities[i], O2.cities[i] = O2.cities[i], O1.cities[i]  # swap segments
    for i in range(len(O1.cities)):
        if i < a or i > b:
            while O1.cities[i] in O1.cities[a:b + 1]:
                repeat_index = O1.cities[a:b + 1].index(O1.cities[i])
                O1.cities[i] = O2.cities[a:b + 1][repeat_index]  # replace i with allele in O2
            while O2.cities[i] in O2.cities[a:b + 1]:
                repeat_index = O2.cities[a:b + 1].index(O2.cities[i])
                O2.cities[i] = O1.cities[a:b + 1][repeat_index]
    O1.salesman_correction()
    O2.salesman_correction()
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def order_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = list(range(0, len(P1.cities)))
    a, b = sorted(sample(indexes, 2))
    seq1, seq2 = [], []
    for i in list(range(b + 1, len(P1.cities))) + list(range(0, b + 1)):
        if P2.cities[i] not in O1.cities[a:b + 1]:
            seq2.append(P2.cities[i])
        if P1.cities[i] not in O2.cities[a:b + 1]:
            seq1.append(P1.cities[i])
    j = 0
    for i in list(range(b + 1, len(P1.cities))) + list(range(0, a)):
        O1.cities[i] = seq2[j]
        O2.cities[i] = seq1[j]
        j = j + 1
    O1.salesman_correction()
    O2.salesman_correction()
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def cycle_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = [_ for _ in range(len(P1.cities))]
    first = sample(list(range(len(P1.cities))), k=1)[0]  # choose the first city in P1.cities to inherit
    i = first
    while P2.cities[i] != O1.cities[first]:
        indexes[i] = -1  # record the inherited index of P1.cities
        i = P1.cities.index(P2.cities[i])  # index of the allele in P1.cities
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O1.cities[i] = P2.cities[i]
    indexes = [_ for _ in range(len(P2.cities))]
    first = sample(list(range(len(P2.cities))), k=1)[0]  # choose the first city in P1.cities to inherit
    i = first
    while P1.cities[i] != O2.cities[first]:
        indexes[i] = -1  # record the inherited index of P1.cities
        i = P2.cities.index(P1.cities[i])  # index of the allele in P1.cities
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O2.cities[i] = P1.cities[i]
    O1.salesman_correction()
    O2.salesman_correction()
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def cycle_crossover_2(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = [_ for _ in range(len(P1.cities))]
    first = sample(list(range(len(P1.cities))), k=1)[0]  # choose the first city in P1.cities to inherit
    i = first
    while P2.cities[i] != O1.cities[first]:
        indexes[i] = -1  # record the inherited index of P1.cities
        i = P1.cities.index(P2.cities[i])  # index of the allele in P1.cities
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O1.cities[i] = P2.cities[i]
            O1.salesmen[i] = P2.salesmen[i]
    indexes = [_ for _ in range(len(P2.cities))]
    first = sample(list(range(len(P2.cities))), k=1)[0]  # choose the first city in P1.cities to inherit
    i = first
    while P1.cities[i] != O2.cities[first]:
        indexes[i] = -1  # record the inherited index of P1.cities
        i = P2.cities.index(P1.cities[i])  # index of the allele in P1.cities
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O2.cities[i] = P1.cities[i]
            O2.salesmen[i] = P1.salesmen[i]
    O1.salesman_correction()
    O2.salesman_correction()
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def mutate(solution):
    indexes = list(range(0, len(solution.cities)))
    i, j = sorted(sample(indexes, 2))
    solution.cities[i], solution.cities[j] = solution.cities[j], solution.cities[i]
    if A[solution.cities[i]][solution.salesmen[i]] == 0:
        salesman = sample(salesmen_of(solution.cities[i]), 1)[0]
        solution.salesmen[i] = salesman
    if A[solution.cities[j]][solution.salesmen[j]] == 0:
        salesman = sample(salesmen_of(solution.cities[j]), 1)[0]
        solution.salesmen[j] = salesman
    solution.calculate_length()
    solution.validate()
    return solution


def _2_OPT(route):
    improve = 0
    for i in range(1, len(route)):
        for j in range(i + 1, len(route)):
            d1 = Distance[route[i - 1]][route[i]] + Distance[route[j - 1]][route[j]]
            d2 = Distance[route[i - 1]][route[j - 1]] + Distance[route[i]][route[j]]
            if d2 < d1:
                route = route[:i] + list(reversed(route[i:j])) + route[j:]
                improve += d2 - d1
    return route, improve


plt.figure(figsize=(10, 4))
max_epoch = 30
max_gen = 10000
max_eva = 1000 * instance.n
NP = 100  # size of population
Pc = 0.5  # cross prob
Pm = 0.1  # mutation prob
PMX = partially_mapped_crossover
OX = order_crossover
CX = cycle_crossover
CX2 = cycle_crossover_2
# ---------> Change save file here! <---------
Crossover = CX2
file_name = 'CX+'
if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}/{m}'
    make_dir(save_path)
    start_time = time()
    optimum = Solution(init='RND')
    for epoch in range(1, max_epoch + 1):
        POP = []
        # Initialization.
        for _ in range(NP):
            POP.append(Solution(init='GRD'))
        # Progress
        pBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest_list = []
        pBest_list = []
        pWorst_list = []
        for gen in range(1, max_gen + 1):
            # Select.
            offspring = []
            fit = np.array([1 / (1 + POP[i].length) for i in range(len(POP))])  # fitness = 1 / (1 + length)
            fit = fit / fit.sum()  # fitness proportion
            winners = np.random.choice(a=np.arange(0, len(POP)), size=NP, replace=True,
                                       p=fit)  # roulette wheel selection
            for winner in winners:  # add winners to offspring population
                offspring.append(POP[winner])
            offspring[offspring.index(max(offspring, key=lambda _: _.length))] = deepcopy(gBest)  # elite strategy
            POP = offspring  # replace the old POP
            # Crossover.
            shuffle(POP)
            for i in range(1, NP, 2):
                if uniform(0, 1) < Pc:
                    POP[i - 1], POP[i] = CX(POP[i - 1], POP[i])
            # Mutation.
            for i in range(1, NP):
                if uniform(0, 1) < Pm:
                    POP[i] = mutate(POP[i])
            # Record gBest
            pBest = deepcopy(min(POP, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            pWorst_list.append(max(POP, key=lambda _: _.length).length)
            # report
            print(f'method=GA, case={case}, m={m}, NP={NP}, Pc={Pc}, Pm={Pm}. '
                  f'At epoch {epoch} gen {gen}: {round(gBest.length, 2)}', end='\r')
            # plot curves & routes
            # plt.clf()
            # plt.subplot(1, 2, 1)
            # plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
            # plt.plot(range(len(pBest_list)), pBest_list, label='pBest')
            # plt.plot(range(len(pWorst_list)), pWorst_list, label='pWorst')
        #     plt.legend()
        #     plt.title(f'total length: {gBest.length:.2f}')
        #     plt.subplot(1, 2, 2)
        #     instance.draw(plt, routes=gBest.to_routes())
        #     plt.pause(0.001)
        # plt.show()
        if optimum.length > gBest.length:
            optimum = deepcopy(gBest)
        # print(optimum.to_routes(single=True))
        # save info per epoch
        mode = 'w' if epoch == 1 else 'ab'
        with open(save_path + f'/fitness_{file_name}.txt', mode) as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))
        # end report
        print(f'method=GA, case={case}, m={m}, NP={NP}, Pc={Pc}, Pm={Pm}. '
              f'At epoch {epoch} gen {max_gen}: {round(gBest.length, 2)}')
    # save optimal route
    route = np.array(optimum.to_routes(single=True)).astype(int)
    with open(save_path + f'/route_{file_name}.txt', 'w') as f:
        np.savetxt(f, route.reshape(1, len(route)), fmt='%d')
    # print time cost
    print(f"Time cost {strftime('%H:%M:%S', gmtime(time() - start_time))}")
